import pylab;

class DataModel(object):
    def __init__(self):
        self.nActions = 4;
        self.nStates = 12;
        self.nTimeEvents = 3;
        self.nStActs = self.nActions * self.nStates;
        self.gralMdpParams = [self.nStates,self.nActions,self.nTimeEvents,0];
        self.stDimNames = ["Type-Walls", "X", "Y", "Energy", "Time"];
        self.stRiskTypes = ["State", "Energy", "Time"];
        self.nRisks = len(self.stRiskTypes);
        self.nDims = len(self.stDimNames);
        self.stateDims = pylab.zeros((self.nStates,len(self.stDimNames)));
        self.mdpVals = pylab.zeros((self.nStates,3));
        self.qVals = pylab.zeros((self.nStates,self.nActions));
        self.stateTrans = pylab.zeros((self.nActions,self.nStates,self.nStates));
        self.stateDists = pylab.zeros((self.nStates,self.nStates));
        self.selectedDists = [0.0 for i in range(self.nStates)]; #@UnusedVariable
        self.nStInfoCols = self.nDims + 3 + self.nActions + 1 + self.nRisks*3;
        self.nRiskActResCols = 10;
        self.nMdpRiskCols = 4 + self.nStInfoCols + self.nRiskActResCols-2;
        self.stInfo = self.initStInfo();
        self.riskActRes, self.riskActResBs = self.initRiskActRes();
        self.stTraceEvents = [[0 for i in range(3)] for j in range(self.nTimeEvents)]; #@UnusedVariable
        self.distTraceLog, self.distTrace, self.distTraceRow = [], [], [];
        self.nDistTimeEvents = 0;
        self.nMaxD = 0;
#        self.multiCriteriaR = [];
#        self.stateRiskLimits = [];
#        self.stateTransAlt = [];
#        self.modelingAlgorithmVals = [];
#        self.stateValLearningVals = [];
#        self.explorationExploitMethods = [];
#        self.graphSettings = [];
#        self.gridSettings = []; 
        self.initMaxMins();
        self.dfGrpActs = None;
        
    def initMaxMins(self):
        self.nMaxQ = [0 for i in range(self.nStates)]; #@UnusedVariable
        self.nMinQ = [0 for i in range(self.nStates)]; #@UnusedVariable
        self.nMinR, self.nMaxR = 0, 0;
        self.nMinV, self.nMaxV = 0, 0;
        
    def getDataTables(self):
        aTables = [];
        aTables.append(self.nActions);
        aTables.append(self.nStates);
        aTables.append(self.nTimeEvents);
        aTables.append(self.nStActs);
        aTables.append(self.gralMdpParams);
        aTables.append(self.stDimNames);
        aTables.append(self.stRiskTypes);
        aTables.append(self.nRisks);
        aTables.append(self.nDims);
        aTables.append(self.stateDims);
        aTables.append(self.mdpVals);
        aTables.append(self.qVals);
        aTables.append(self.stateTrans);
        aTables.append(self.stateDists);
        aTables.append(self.selectedDists);
        aTables.append(self.nStInfoCols);
        aTables.append(self.nRiskActResCols);
        aTables.append(self.nMdpRiskCols);
        aTables.append(self.stInfo);
        aTables.append(self.riskActRes);
        aTables.append(self.riskActResBs);
        aTables.append(self.stTraceEvents);        
        aTables.append(self.distTraceLog);
        aTables.append(self.distTrace);
        aTables.append(self.distTraceRow);
        aTables.append(self.nDistTimeEvents);
        aTables.append(self.nMaxD);
        return aTables;
        
    def setDataTables(self, aTables):
        self.nActions = aTables[0];
        self.nStates = aTables[1];
        self.nTimeEvents = aTables[2];
        self.nStActs = aTables[3];
        self.gralMdpParams = aTables[4];
        self.stDimNames = aTables[5];
        self.stRiskTypes = aTables[6];
        self.nRisks = aTables[7];
        self.nDims = aTables[8];
        self.stateDims = aTables[9];
        self.mdpVals = aTables[10];
        self.qVals = aTables[11];
        self.stateTrans = aTables[12];
        self.stateDists = aTables[13];
        self.selectedDists = aTables[14];
        self.nStInfoCols = aTables[15];
        self.nRiskActResCols = aTables[16];
        self.nMdpRiskCols = aTables[17];
        self.stInfo = aTables[18];
        self.riskActRes = aTables[19];
        self.riskActResBs = aTables[20];
        self.stTraceEvents = aTables[21];        
        self.distTraceLog = aTables[22];        
        self.distTrace = aTables[23];        
        self.distTraceRow = aTables[24];        
        self.nDistTimeEvents = aTables[25];        
        self.nMaxD = aTables[26];
        
    def initStTrans(self):        
        nStates = self.nStates;
        nActions = self.nActions;
        self.stateTrans = pylab.zeros((nActions,nStates,nStates));
        
    def updateStTrans(self):
        nStates = self.nStates;
        nActions = self.nActions;
        newTrans = pylab.zeros((nActions,nStates,nStates));
        old_nActions = pylab.size(self.stateTrans,0);
        old_nStates = pylab.size(self.stateTrans,1);
        for i in range(old_nActions):
            for j in range(old_nStates):
                for k in range(old_nStates):
                    if i < nActions and j < nStates and k < nStates:
                        newTrans[i][j][k] = self.stateTrans[i][j][k];
        self.stateTrans = newTrans.copy();
        
    def updateStDists(self):
        nStates = self.nStates;
        self.selectedDists = [0.0 for i in range(nStates)]; #@UnusedVariable
        newDists = pylab.zeros((nStates,nStates));
        old_nStates = len(self.stateDists);
        for i in range(old_nStates):
            for j in range(old_nStates):
                if i < nStates and j < nStates:
                    newDists[i][j] = self.stateDists[i][j];
        self.stateDists = newDists.copy();
        
    def updateStInfo(self):
        nStates = self.nStates;
        nActions = self.nActions; #@UnusedVariable
        nRows = len(self.stInfo)
        if nRows < nStates:
            for i in range(nStates-nRows): #@UnusedVariable
                self.stInfo.append(self.getStInfoRow()); #@UnusedVariable
        elif nRows > nStates:
            for i in range(nRows-nStates): #@UnusedVariable
                self.stInfo.remove(self.stInfo[len(self.stInfo)-1]); #@UnusedVariable
                
    def initStInfo(self):
        stInfo = [];
        for i in range(self.nStates): #@UnusedVariable
            stInfo.append(self.getStInfoRow());
        return stInfo;
    
    def getStInfoRow(self):
        aRow = [];
        for i in range(self.nDims): aRow.append(0); #@UnusedVariable
        aRow.append(0.0); #Reward
        aRow.append(0.0); #Value
        aRow.append(0); #Policy
        for i in range(self.nActions): aRow.append(0.0); #@UnusedVariable
        aRow.append("Normal"); #State Type
        for i in range(self.nRisks): aRow.append("No"); #@UnusedVariable
        for i in range(self.nRisks*2): aRow.append(0.0); #@UnusedVariable
        return aRow;
    
    def initRiskActRes(self):
        aRiskAct = [];
        aRiskActBs = [];
        for i in range(self.nStates):
            for j in range(self.nActions):
                aRiskAct.append(self.getRiskActResRow(i+1, j+1));
                aRiskActBs.append([True, True, True]);
        return aRiskAct, aRiskActBs;
    
    def getRiskActResRow(self, iSt, iAct):
        res = "Good Action";
        st = "State " + str(iSt);
        act = "Action " + str(iAct);
        return [st,act, 0.0,">=",0.0, "&", 0.0,"<=",0.0, res];
    
    def updateRiskActRes(self):
        self.nStActs = self.nActions * self.nStates;
        nRows = len(self.riskActRes);
        nStates = self.nStates;
        nActions = self.nActions;
        nOldSt = nRows/nActions;
        nOldAct = nRows/nStates; #@UnusedVariable
        if nOldSt < nStates: #Add States
            for i in range(nStates-nOldSt):
                for j in range(nActions):
                    iSt = nOldSt+(i+1);
                    self.riskActRes.append(self.getRiskActResRow(iSt, j+1));
                    self.riskActResBs.append([True, True, True]);
        elif nOldSt > nStates: #Remove States
            for i in range((nOldSt-nStates)*nActions): #@UnusedVariable
                self.riskActRes.remove(self.riskActRes[len(self.riskActRes)-1]);
                self.riskActResBs.remove(self.riskActResBs[len(self.riskActResBs)-1]);
        
    def updateStTraceEvents(self):
        nTimeEvents = self.nTimeEvents;
        nRows = len(self.stTraceEvents)
        if nRows < nTimeEvents:
            for i in range(nTimeEvents-nRows): #@UnusedVariable
                self.stTraceEvents.append([0,0,0]);
        elif nRows > nTimeEvents:
            for i in range(nRows-nTimeEvents): #@UnusedVariable
                lastRow = self.stTraceEvents[len(self.stTraceEvents)-1];
                self.stTraceEvents.remove(lastRow); #@UnusedVariable

    
